#!/usr/bin/env python3

"""
性能监控脚本 - 用于监控无人机导航系统的性能
专为Jetson Orin Nano优化
"""

import rclpy
from rclpy.node import Node
from std_msgs.msg import String
import psutil
import json
import time
from datetime import datetime
import os
import threading
import subprocess

class PerformanceMonitor(Node):
    
    def __init__(self):
        super().__init__('performance_monitor')
        
        # 发布性能数据
        self.performance_pub = self.create_publisher(
            String, 
            '/system/performance', 
            10
        )
        
        # 参数声明
        self.declare_parameter('monitor_interval', 2.0)  # 监控间隔(秒)
        self.declare_parameter('log_to_file', True)      # 是否记录到文件
        self.declare_parameter('warning_cpu_threshold', 80.0)    # CPU警告阈值
        self.declare_parameter('warning_memory_threshold', 85.0) # 内存警告阈值
        self.declare_parameter('warning_temp_threshold', 75.0)   # 温度警告阈值
        
        # 获取参数
        self.monitor_interval = self.get_parameter('monitor_interval').value
        self.log_to_file = self.get_parameter('log_to_file').value
        self.cpu_threshold = self.get_parameter('warning_cpu_threshold').value
        self.memory_threshold = self.get_parameter('warning_memory_threshold').value
        self.temp_threshold = self.get_parameter('warning_temp_threshold').value
        
        # 初始化日志文件
        if self.log_to_file:
            self.log_dir = os.path.expanduser('~/drone_navigation_ws/logs')
            os.makedirs(self.log_dir, exist_ok=True)
            self.log_file = os.path.join(
                self.log_dir, 
                f'performance_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log'
            )
        
        # 创建定时器
        self.timer = self.create_timer(
            self.monitor_interval, 
            self.monitor_performance
        )
        
        # 警告状态跟踪
        self.warning_states = {
            'cpu': False,
            'memory': False,
            'temperature': False
        }
        
        self.get_logger().info("性能监控器已启动")
        if self.log_to_file:
            self.get_logger().info(f"性能日志保存到: {self.log_file}")
    
    def get_cpu_info(self):
        """获取CPU信息"""
        try:
            cpu_percent = psutil.cpu_percent(interval=1)
            cpu_count = psutil.cpu_count()
            cpu_freq = psutil.cpu_freq()
            
            return {
                'usage_percent': cpu_percent,
                'core_count': cpu_count,
                'frequency_mhz': cpu_freq.current if cpu_freq else 0,
                'max_frequency_mhz': cpu_freq.max if cpu_freq else 0
            }
        except Exception as e:
            self.get_logger().error(f"获取CPU信息失败: {e}")
            return {}
    
    def get_memory_info(self):
        """获取内存信息"""
        try:
            memory = psutil.virtual_memory()
            swap = psutil.swap_memory()
            
            return {
                'total_gb': round(memory.total / (1024**3), 2),
                'used_gb': round(memory.used / (1024**3), 2),
                'free_gb': round(memory.free / (1024**3), 2),
                'usage_percent': memory.percent,
                'swap_usage_percent': swap.percent,
                'swap_used_gb': round(swap.used / (1024**3), 2)
            }
        except Exception as e:
            self.get_logger().error(f"获取内存信息失败: {e}")
            return {}
    
    def get_disk_info(self):
        """获取磁盘信息"""
        try:
            disk = psutil.disk_usage('/')
            return {
                'total_gb': round(disk.total / (1024**3), 2),
                'used_gb': round(disk.used / (1024**3), 2),
                'free_gb': round(disk.free / (1024**3), 2),
                'usage_percent': round((disk.used / disk.total) * 100, 2)
            }
        except Exception as e:
            self.get_logger().error(f"获取磁盘信息失败: {e}")
            return {}
    
    def get_network_info(self):
        """获取网络信息"""
        try:
            net_io = psutil.net_io_counters()
            return {
                'bytes_sent_mb': round(net_io.bytes_sent / (1024**2), 2),
                'bytes_recv_mb': round(net_io.bytes_recv / (1024**2), 2),
                'packets_sent': net_io.packets_sent,
                'packets_recv': net_io.packets_recv
            }
        except Exception as e:
            self.get_logger().error(f"获取网络信息失败: {e}")
            return {}
    
    def get_temperature_info(self):
        """获取温度信息（Jetson特有）"""
        try:
            temps = {}
            
            # 尝试读取Jetson温度传感器
            temp_paths = [
                '/sys/class/thermal/thermal_zone0/temp',
                '/sys/class/thermal/thermal_zone1/temp',
                '/sys/class/thermal/thermal_zone2/temp'
            ]
            
            for i, path in enumerate(temp_paths):
                try:
                    if os.path.exists(path):
                        with open(path, 'r') as f:
                            temp_raw = f.read().strip()
                            temp_celsius = int(temp_raw) / 1000.0
                            temps[f'zone_{i}'] = round(temp_celsius, 1)
                except:
                    continue
            
            # 尝试使用psutil获取温度
            if hasattr(psutil, 'sensors_temperatures'):
                sensors = psutil.sensors_temperatures()
                for name, entries in sensors.items():
                    for entry in entries:
                        temps[f'{name}_{entry.label or "temp"}'] = entry.current
            
            return temps
        except Exception as e:
            self.get_logger().error(f"获取温度信息失败: {e}")
            return {}
    
    def get_gpu_info(self):
        """获取GPU信息（如果可用）"""
        try:
            # 尝试使用nvidia-smi获取GPU信息
            result = subprocess.run(
                ['nvidia-smi', '--query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu', 
                 '--format=csv,noheader,nounits'],
                capture_output=True, text=True, timeout=5
            )
            
            if result.returncode == 0:
                lines = result.stdout.strip().split('\n')
                gpu_info = {}
                for i, line in enumerate(lines):
                    parts = line.split(', ')
                    if len(parts) == 4:
                        gpu_info[f'gpu_{i}'] = {
                            'utilization_percent': float(parts[0]),
                            'memory_used_mb': float(parts[1]),
                            'memory_total_mb': float(parts[2]),
                            'temperature_celsius': float(parts[3])
                        }
                return gpu_info
            
        except Exception as e:
            # GPU信息获取失败不是关键错误
            pass
        
        return {}
    
    def get_process_info(self):
        """获取关键进程信息"""
        try:
            processes = []
            
            # 查找ROS2相关进程
            for proc in psutil.process_iter(['pid', 'name', 'cpu_percent', 'memory_percent']):
                try:
                    if any(keyword in proc.info['name'].lower() for keyword in 
                           ['ros', 'navigation', 'point_lio', 'livox']):
                        processes.append({
                            'pid': proc.info['pid'],
                            'name': proc.info['name'],
                            'cpu_percent': proc.info['cpu_percent'],
                            'memory_percent': round(proc.info['memory_percent'], 2)
                        })
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue
            
            return processes[:10]  # 只返回前10个进程
        except Exception as e:
            self.get_logger().error(f"获取进程信息失败: {e}")
            return []
    
    def check_warnings(self, performance_data):
        """检查警告条件"""
        warnings = []
        
        # CPU警告
        cpu_usage = performance_data.get('cpu', {}).get('usage_percent', 0)
        if cpu_usage > self.cpu_threshold:
            if not self.warning_states['cpu']:
                warnings.append(f"CPU使用率过高: {cpu_usage:.1f}%")
                self.warning_states['cpu'] = True
        else:
            if self.warning_states['cpu']:
                warnings.append(f"CPU使用率恢复正常: {cpu_usage:.1f}%")
                self.warning_states['cpu'] = False
        
        # 内存警告
        memory_usage = performance_data.get('memory', {}).get('usage_percent', 0)
        if memory_usage > self.memory_threshold:
            if not self.warning_states['memory']:
                warnings.append(f"内存使用率过高: {memory_usage:.1f}%")
                self.warning_states['memory'] = True
        else:
            if self.warning_states['memory']:
                warnings.append(f"内存使用率恢复正常: {memory_usage:.1f}%")
                self.warning_states['memory'] = False
        
        # 温度警告
        temperatures = performance_data.get('temperature', {})
        max_temp = max(temperatures.values()) if temperatures else 0
        if max_temp > self.temp_threshold:
            if not self.warning_states['temperature']:
                warnings.append(f"温度过高: {max_temp:.1f}°C")
                self.warning_states['temperature'] = True
        else:
            if self.warning_states['temperature'] and max_temp > 0:
                warnings.append(f"温度恢复正常: {max_temp:.1f}°C")
                self.warning_states['temperature'] = False
        
        return warnings
    
    def monitor_performance(self):
        """主监控函数"""
        try:
            # 收集性能数据
            performance_data = {
                'timestamp': datetime.now().isoformat(),
                'cpu': self.get_cpu_info(),
                'memory': self.get_memory_info(),
                'disk': self.get_disk_info(),
                'network': self.get_network_info(),
                'temperature': self.get_temperature_info(),
                'gpu': self.get_gpu_info(),
                'processes': self.get_process_info()
            }
            
            # 检查警告
            warnings = self.check_warnings(performance_data)
            if warnings:
                performance_data['warnings'] = warnings
                for warning in warnings:
                    self.get_logger().warning(warning)
            
            # 发布性能数据
            msg = String()
            msg.data = json.dumps(performance_data, ensure_ascii=False)
            self.performance_pub.publish(msg)
            
            # 记录到文件
            if self.log_to_file:
                with open(self.log_file, 'a', encoding='utf-8') as f:
                    f.write(f"{json.dumps(performance_data, ensure_ascii=False)}\n")
            
            # 定期打印简要信息
            if hasattr(self, 'last_summary_time'):
                if time.time() - self.last_summary_time > 30:  # 每30秒打印一次
                    self.print_summary(performance_data)
                    self.last_summary_time = time.time()
            else:
                self.last_summary_time = time.time()
                self.print_summary(performance_data)
            
        except Exception as e:
            self.get_logger().error(f"性能监控出错: {e}")
    
    def print_summary(self, data):
        """打印性能摘要"""
        cpu = data.get('cpu', {})
        memory = data.get('memory', {})
        temps = data.get('temperature', {})
        
        summary = f"性能摘要 - CPU: {cpu.get('usage_percent', 0):.1f}%, "
        summary += f"内存: {memory.get('usage_percent', 0):.1f}%, "
        if temps:
            max_temp = max(temps.values())
            summary += f"最高温度: {max_temp:.1f}°C"
        
        self.get_logger().info(summary)

def main():
    rclpy.init()
    
    try:
        monitor = PerformanceMonitor()
        rclpy.spin(monitor)
    except KeyboardInterrupt:
        pass
    except Exception as e:
        print(f"监控器异常: {e}")
    finally:
        rclpy.shutdown()

if __name__ == '__main__':
    main() 